import sys

sys.path.append(r'D:\vc\Rb87Python')
sys.path.append(r'D:\vc\Rb87Python\theory3D_SOC\python')
import xlrd
import numpy as np
# import scipy.io as sio
from matplotlib import pyplot as plt
import matplotlib.patches as patches
import os
import time
import collections
import itertools

import HotAtomDist_3D
import Spin_Texture_Initial_Settings

import basic_function.Rb87BEC_Basic_Functions as bf
import basic_function.Rb87BEC_Fit as bffit

import socket

os.system('chcp 65001')
Host_Name = socket.gethostname()
if Host_Name == 'VictorChengPC':
    Computer_Name = 'X:\\Work\\mail.ustc.edu.cn\\Rb87BEC - Files\\'
elif Host_Name == 'calcenter2':
    Computer_Name = 'D:\\vc\\mail.ustc.edu.cn\\Rb87BEC - Files\\'

Data_Exl = xlrd.open_workbook(Computer_Name + r'Rb87files\3DSOC_spin_texture\spin_texture.xlsx')
Table = Data_Exl.sheets()[0]
# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
Mat_Data = np.array(range(754, 793 + 1, 1))
# Mat_Data = np.array(range(627, 676 + 1, 1))
# Mat_Data = np.concatenate((np.array([721, 722]), Mat_Data))
# mz_simed = -0.25
Atom_Temperature = 150e-9
Fit_Prefix = ['NoFit', '']
Fit_BEC_Switch = 1
for Row_Num in range(len(Mat_Data)):
    Exl_RowNum = Mat_Data[Row_Num] - 1
    Input_Matrix = Table.row_values(Exl_RowNum)

    YearMonthDay = str(int(Input_Matrix[0]))
    Ini_Photo = int(Input_Matrix[1])
    Fin_Photo = int(Input_Matrix[2])
    Data_Folder_Name = Computer_Name + 'Rb87Data\\3DSOC_spin_texture\\' + \
        Fit_Prefix[Fit_BEC_Switch] + '\\' + YearMonthDay + '\\All_Spin_Polarization\\'
    Photo = str(Ini_Photo) + 'To' + str(Fin_Photo)
    Cam = Input_Matrix[8]
    mz_exp = Input_Matrix[10]/2
    Exp_Data = np.load(Computer_Name + "Rb87Data\\3DSOC_spin_texture\\" +
                       YearMonthDay + "\\All_Spin_Polarization\\Cam_" +
                       Cam + "_" + Photo + "Spin_Polarization_FBZ.npz", allow_pickle=True)
    Exp_Data_Dict = dict(Exp_Data)
    SPs = Exp_Data_Dict['SPs']
    if 'FBZ_Ups' in Exp_Data_Dict:
        FBZ_Ups = Exp_Data_Dict['FBZ_Ups'].item()
        FBZ_Dns = Exp_Data_Dict['FBZ_Dns'].item()
    else:
        print('Warning, no fbz up and down')
        FBZ_Ups = 0
        FBZ_Dns = 0
    FBZ_Up_Lowers = collections.OrderedDict()
    FBZ_Dn_Lowers = collections.OrderedDict()
    SP_Lowers = collections.OrderedDict()
    try:
        Numerical_Data = np.load(Computer_Name + "Rb87Data\\3DSOCbandData\\191107\\3x3x3;4x66x66x66;"
                                                 "V_xz=1.770E_r,V_y=-3.540E_r,Oxy=1.022E_r,Ozy=-1.022E_r,"
                                                 "T=1.000e-07K,n=3.000e+18m-3del_psi=1.000pi,"
                                                 "mz=%.3fErno_states_v1.npz" % mz_exp)
        energies = Numerical_Data['energies']
        ESigmaZ = Numerical_Data['ESigmaZ']
        (mu, ns, nsp) = HotAtomDist_3D.ChemicalPotential(energies, Atom_Temperature, 3e18)
        for Photo_Index in FBZ_Ups.keys():
            FBZ_Up = FBZ_Ups[Photo_Index]
            FBZ_Dn = FBZ_Dns[Photo_Index]
            n_Up = ns * (1 + ESigmaZ) / 2
            n_Dn = ns * (1 - ESigmaZ) / 2
            Band_Ratio_Up = np.sum(n_Up, axis=2) / np.sum(n_Up, axis=(0, 2))
            Band_Ratio_Dn = np.sum(n_Dn, axis=2) / np.sum(n_Dn, axis=(0, 2))
            FBZ_Up_Lower = FBZ_Up * Band_Ratio_Up[0, ...]
            FBZ_Dn_Lower = FBZ_Dn * Band_Ratio_Dn[0, ...]
            SP_Lower = (FBZ_Up_Lower - FBZ_Dn_Lower) / (FBZ_Up_Lower + FBZ_Dn_Lower)
            FBZ_Up_Lowers[Photo_Index] = FBZ_Up_Lower
            FBZ_Dn_Lowers[Photo_Index] = FBZ_Dn_Lower
            SP_Lowers[Photo_Index] = SP_Lower
        Exp_Data_Dict['FBZ_Up_Lowers'] = FBZ_Up_Lowers
        Exp_Data_Dict['FBZ_Dn_Lowers'] = FBZ_Dn_Lowers
        Exp_Data_Dict['SP_Lowers'] = SP_Lowers
    except FileNotFoundError:
        print('Waring, no band data')
        SP_Lowers = SPs
        Exp_Data_Dict['SP_Lowers'] = SP_Lowers
    np.savez(Computer_Name + "Rb87Data\\3DSOC_spin_texture\\" +
             YearMonthDay + "\\All_Spin_Polarization\\Cam_" +
             Cam + "_" + Photo + "Spin_Polarization_FBZ_WithLower_%.3eK.npz" % Atom_Temperature, **Exp_Data_Dict)
    Photo_Ind = Ini_Photo
    while Photo_Ind <= Fin_Photo:
        Data = np.load(Data_Folder_Name + 'Cam_22860_photo%d_Spin_Polarization_FBZ.npz' % Photo_Ind)
        Data_Dict = dict(Data)
        SP = Data_Dict['SP']
        try:
            te = Band_Ratio_Up
        except NameError:
            print('error')
            SP_Lower = SP
            Data_Dict['SP_Lower'] = SP_Lower
            np.savez(Data_Folder_Name + 'Cam_22860_photo%d_Spin_Polarization_FBZ_WithLower_%.3eK.npz' %
                     (Photo_Ind, Atom_Temperature), **Data_Dict)
            Photo_Ind += 1
            continue
        FBZ_Up = Data_Dict['FBZ_Up']
        FBZ_Dn = Data_Dict['FBZ_Dn']
        FBZ_Up_Lower = FBZ_Up * Band_Ratio_Up[0, ...]
        FBZ_Dn_Lower = FBZ_Dn * Band_Ratio_Dn[0, ...]
        SP_Lower = (FBZ_Up_Lower - FBZ_Dn_Lower) / (FBZ_Up_Lower + FBZ_Dn_Lower)
        Data_Dict['FBZ_Up_Lower'] = FBZ_Up_Lower
        Data_Dict['FBZ_Dn_Lower'] = FBZ_Dn_Lower
        Data_Dict['SP_Lower'] = SP_Lower
        np.savez(Data_Folder_Name + 'Cam_22860_photo%d_Spin_Polarization_FBZ_WithLower_%.3eK.npz' %
                 (Photo_Ind, Atom_Temperature), **Data_Dict)
        Photo_Ind += 1
# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
